#!/usr/bin/python

"""
This file contains a few functions and a script for finding the
eigenvector corresponding to the smallest eigenvalue of a matrix,
which we will call the principal eigenvector for the remainder of this
document. The matrix used will be derived from a convolution kernel 
represented by a real-rational transfer function.
"""

from scipy.signal import lti
from scipy.linalg import sqrtm, inv
from scipy.linalg import inv
#from sympy.mpmath import sqrtm
from scipy.linalg import eig #Full diagonalization. Inefficient.
from scipy.sparse.linalg import eigsh
from scipy.weave import inline, blitz, converters
from numpy import linspace, matrix, ndarray, where

debug = True

def kernel_array(tf, t_f, n_pts):
    """
    Evaluates the impulse response of some transfer function, if it 
    exists.
    Just noticed what I did with the input variable names.  
    """
    if debug:
        print "kernel_array: n_pts = " + str(n_pts)
    if tf is None:
        return None
    return tf.impulse(X0 = 0., T = linspace(0., t_f, n_pts), N=n_pts)[1]

def two_norm_matrix(kernel):
    """
    Evaluates a formula from Discretization.pdf.
    """
    n = len(kernel)
    output = ndarray((n,n),dtype='double')
    if debug:
        print "two_norm_matrix: just before c code"
    code='''
        for (int j = 0; j < n; ++j){
            for (int k = 0; k < n; ++k){
                output(j,k) = 0.;
                for (int l = 0; l < 2*n-1; ++l){
                    if ((l-j >= 0) && (l-j < n) && (l-k >= 0) && (l-k < n)){
                        output(j,k) = output(j,k) + kernel(l-j) * kernel(l-k);
                    }
                }
            }
        }
    '''
    inline(code, ['output','kernel','n'], type_converters = converters.blitz)
    return matrix(output)

def objective_matrix(min_kernel, max_kernel=None, do_max=False):
    """
    Takes one or two one-dimensional arrays which contain the impulse
    responses of the transfer functions that you want to minimize and
    maximize, respectively. Returns a dense numpy matrix whose 
    principal eigenvector is the optimal open-loop control law. 

    do_max controls whether or not the matrix will be constructed for 
    the minimization or equivalent maximization problem. 
    """
    min_mat = two_norm_matrix(min_kernel)
    max_mat = two_norm_matrix(max_kernel)
    
    if debug:
        print "objective_matrix: type of min_mat, max_mat: " + ', '.join(map(lambda a: str(type(a)), (min_mat, max_mat)))
    
    if debug:
        print "objective_matrix: right before sqrtm, inv"
    if do_max:
        outer_mat = sqrtm(inv(min_mat))
        inner_mat = max_mat
    else:
        outer_mat = sqrtm(inv(max_mat))
        inner_mat = min_mat

    if debug:
        print "objective_matrix: type of outer_mat: " + str(type(outer_mat))
        print "objective_matrix: type of inner_mat: " + str(type(inner_mat))
    return outer_mat * inner_mat * outer_mat, outer_mat

def principal_eigenvector(matr, do_max=False):
    """
    Returns the eigenvector corresponding to the smallest eigenvalue
    of the input matrix.
    """
    full=False
    if full:
        #Full Diagonalization
        if debug:
            print "principal_eigenvector: before diagonalization"
        vals, vecs = eig(matr)
        if do_max:
            loc = where(vals == max(vals))[0][0]
        else:
            loc = where(vals == min(vals))[0][0]
        return vecs[:,loc]
    else:
        print matr.dtype
        if do_max:
            return eigsh(matr, k=1)[1]
        else:
            #Weird mode and sigma to make eigsh work, should switch matrices.
            return eigsh(matr, k=1, mode='buckling', sigma=0.001, which='LM')[1]
 
def control_law(min_tf, max_tf = None, t_f=10., n_pts=2000, do_max=False):
    """
    Combines the above functions to obtain an optimal discrete-time 
    control law from one or two `lti` objects.
    """
    min_kernel = kernel_array(min_tf, t_f, n_pts)
    max_kernel = kernel_array(max_tf, t_f, n_pts)
    control_matrix, basis_change = objective_matrix(min_kernel, max_kernel, do_max)
    if debug:
        print "control_law: before return"
    evec = principal_eigenvector(control_matrix, do_max)
    evec = basis_change.dot(evec)
    if debug:
        print "control_law: type of evec: " + str(type(evec))
    return evec

if __name__ == '__main__':
    import cPickle as pickle
    from sys import argv
    t_f = float(argv[1])
    #Transfer function for unwanted signal:
    min_tf = lti([-16., -64., -320., -640., -1024., -768., 0.],
               [1., 8., 56., 224., 672., 1344., 1856., 1536., 576.])
    #Transfer function of desired signal:
    max_tf = lti([-96., 0., 768., 2304., 2304.],
               [1., 8., 56., 224., 672., 1344., 1856., 1536., 576.])
    
    n_pts = 2000

    max_ratio_ctrl = control_law(min_tf, max_tf, t_f, n_pts=n_pts, do_max=True)
    
    with open('_'.join(['saved_pulses/pulse', str(int(t_f)), str(n_pts), '.pkl']),'w') as phil:
        pickle.dump(max_ratio_ctrl, phil)